import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report, accuracy_score
import torch
from torch import nn
import torch.utils.data as Data
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
from torchvision.datasets import CIFAR10
from torchvision.utils import make_grid
from torchinfo import summary
from tqdm import tqdm
from ae_net import AE_FC, AE_CNN, VAE_FC, VAE_CNN
import torchvision
import torchvision.datasets as dset
from train import train_model, train_vae_model

def train_AE_FC():
    #生成模型
    image_size=32
    net = AE_FC(image_size, 64)
    summary(net, (1, 3, image_size, image_size))

    trans = transforms.Compose((transforms.Resize(image_size), transforms.ToTensor()))
    # 要对从cifar中下载的图像进行一些转变，所以这里先初始化transform
    back_size = 1
    batch_size = 256

    dataset_name = "cifar10"
    trainset = CIFAR10(
        root=r'D:/Data',
        train=True,
        download=False,
        transform=trans)  # 训练数据集

    optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)  # 选好优化方式

    loss_fn = nn.MSELoss(reduction='sum')  # 损失函数为交叉熵，多用于多分类问题 reduction='sum' mean

    train_model(net, dataset_name, trainset, optimizer,loss_fn, back_size, batch_size)

def trian_AE_CNN():
    #生成模型
    image_size=32
    net = AE_CNN()
    summary(net, (1, 3, image_size, image_size))

    trans = transforms.Compose((transforms.Resize(image_size), transforms.ToTensor()))
    # 要对从cifar中下载的图像进行一些转变，所以这里先初始化transform
    back_size = 1
    batch_size = 256

    dataset_name = "cifar10"
    trainset = CIFAR10(
        root=r'D:/Data',
        train=True,
        download=False,
        transform=trans)  # 训练数据集

    optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)  # 选好优化方式

    loss_fn = nn.MSELoss(reduction='sum')  # 损失函数为交叉熵，多用于多分类问题 reduction='sum' mean

    train_model(net, dataset_name, trainset, optimizer,loss_fn, back_size, batch_size)

def train_VAE_FC():
    #生成模型
    image_size=32
    net = VAE_FC(image_size, 64)
    summary(net, (1, 3, image_size, image_size))

    trans = transforms.Compose((transforms.Resize(image_size), transforms.ToTensor()))
    # 要对从cifar中下载的图像进行一些转变，所以这里先初始化transform
    back_size = 1
    batch_size = 1600

    dataset_name = "cifar10"
    trainset = CIFAR10(
        root=r'D:/Data',
        train=True,
        download=False,
        transform=trans)  # 训练数据集

    optimizer = torch.optim.Adam(net.parameters(), lr=1e-4)  # 选好优化方式
    # optimizer = torch.optim.SGD(net.parameters(), lr=1e-3)  # 选好优化方式

    loss_fn = nn.MSELoss(reduction='sum')  # 损失函数为交叉熵，多用于多分类问题 reduction='sum' mean

    train_vae_model(net, dataset_name, trainset, optimizer,loss_fn, back_size, batch_size)

def train_VAE_CNN():
    #生成模型
    image_size=32
    net = VAE_CNN(image_size,)
    summary(net, (1, 3, image_size, image_size))

    trans = transforms.Compose((transforms.Resize(image_size), transforms.ToTensor()))
    # 要对从cifar中下载的图像进行一些转变，所以这里先初始化transform
    back_size = 1
    batch_size = 4096

    dataset_name = "cifar10"
    trainset = CIFAR10(
        root=r'D:/Data',
        train=True,
        download=False,
        transform=trans)  # 训练数据集

    # optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)  # 选好优化方式
    # optimizer = torch.optim.SGD(net.parameters(), lr=1e-5)  # 选好优化方式
    optimizer = torch.optim.RMSprop(net.parameters(), lr=1e-5)  # 选好优化方式

    loss_fn = nn.MSELoss(reduction='sum')  # 损失函数为交叉熵，多用于多分类问题 reduction='sum' mean

    train_vae_model(net, dataset_name, trainset, optimizer,loss_fn, back_size, batch_size)

if __name__ == '__main__':
    # train_VAE_FC()
    train_VAE_CNN()